#/usr/bin/python
# -*- encoding=UTF-8 -*-

"""class for linux

Class:
    BaseLinux

Author:
    Limbo
"""

import re
import os
import paramiko  # for ssh
import socket
import time


class BaseLinux():
    def __init__(self, name, info, recorder=None, logger=None):
        self.name = name
        self.info = info
        self.logger = logger
        self.recorder = recorder
        self.ssh = paramiko.SSHClient() 
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        #paramiko.util.log_to_file('paramiko.log')
        self.connected = False
        self.root_password = None
        self.stdout = None
        self.stderr = None
        self._cwd = None
        self.timestamp = 0
        self.set_prompt()

    def __del__(self):
        self.close()

    def close(self):
        """close the ssh connect"""
        try:
            self.ssh.close()
        finally:
            self.connected = False

    def log(self, level, msg, **kwargs):
        """recall the self.logger(which is the calss RecordLogger), 
        and add the head parameter.
        @level: values(notice|debug|info|warn|error|exception)
        @kwargs: as end='\n'
        """
        if self.logger and hasattr(self.logger, level):
            if not kwargs.has_key('head'):
                kwargs['head'] = self.name
            getattr(self.logger, level)(msg, **kwargs)
        # to delete
        else:
            print "self.logger has not attr %s" % level

    def record(self, msg, func='write', **kwargs):
        """re-call self.recorder(which is the calss RecordFile)
        @func: value(write|write_file|record)
               the call function for self.recorder, 
        @kwargs: just end='\n' for call record
        """
        if self.recorder and hasattr(self.recorder, func):
            getattr(self.recorder, func)(msg+'\n', **kwargs)
        # to delete
        else:
            print "self.recorder has not fucntion %s" % func

    def connect(self, timeout=10):
        objects = self.info.__dict__
        objects['port'] = 22
        objects['timeout'] = timeout
        #objects['password'] = str(objects['password'])
        
        msg = "connect to %s[%s]" % (self.name, self.info.hostname)
        self.log("info", msg, end=' ')
        try:
            self.ssh.connect(**objects)
        except Exception:
            self.log("notice", "failed", head='', tag='r')
            self.log("exception", "failed ssh connect %s" % self.info.hostname)
            self.connected = False
        else:
            self.log("notice", "successful", head='', tag='g')
            self.connected = True
            self.timestamp = int(time.time())
        return self.connected

    def set_root_password(self, password):
        self.root_password = password

    def set_prompt(self, prompt="$ "):
        self.prompt = prompt

    def get_timestamp(self):
        return self.timestamp

    def chdir(self, path=None):
        self._cwd = path

    def sudo(self, cmd, timeout=30.0, cd=False):
        self.send("whoami")
        if self.stdout == "root":
            self.send(cmd, timeout=timeout, cd=cd)
            return True
        else:
            if not self.root_password:
                self.stderr = "su root error due to no password for %s." % cmd
                self.log("error", self.stderr)
                return False
            invoke = self.ssh.invoke_shell()
            invoke.settimeout(timeout)
            time.sleep(0.2)
            su_cmd = "su root"
            invoke.send(su_cmd + "\n")
            self.record(self.prompt + su_cmd)
            buff = ""
            while not buff.endswith("Password: "):
                resp = invoke.recv(512)
                buff += resp
            invoke.send(self.root_password + "\n")
            
            buff = ""
            while not buff.endswith("# ") and not buff.endswith("$ "):
                resp = invoke.recv(512)
                buff += resp
            if "#" not in buff:
                self.stderr = "su login error " + buff.splitlines()[0]
                self.stdout = ""
                self.log("error", self.stderr)
                return False
            self.prompt = "# "

            if cd and self._cwd :
                new_cmd = "cd {0} && {1}".format(self._cwd, cmd)
            else:
                new_cmd = cmd

            self.record(self.prompt + new_cmd)
            invoke.send(new_cmd + "\n")

            buff = ""
            while not buff.endswith("# "):
                resp = invoke.recv(512)
                buff += resp
            #remove command line
            buff_cmd=""
            output_lines=[]
            for line in buff.splitlines():
                if buff_cmd == new_cmd:
                    output_lines.append(line)
                else:
                    buff_cmd += line.strip()

            self.stdout = '\n'.join(output_lines[:-1])
            self.stderr = ""

            self.record(self.stdout)
            self.prompt = "$ "

            return True

    def send(self, cmd, timeout=30.0, cd=False, revert=True):
        if not self.connected:
            return None
        self.timestamp = int(time.time())
        if cd and self._cwd :
            new_cmd = "cd {0} && {1}".format(self._cwd, cmd)
        else:
            new_cmd = cmd

        self.record(self.prompt + new_cmd)

        try:
            stdout, stderr = \
                self.ssh.exec_command(new_cmd, timeout=timeout)[1:]
        except socket.timeout:
            self.stdout = ""
            self.stderr = ""
        except Exception:
            self.log("exception", "send command [%s] exception" % new_cmd)
            self.close()
            self.stdout = ""
            self.stderr = ""
            if revert and self.connect():
                return self.send(cmd, timeout, cd, False)
        else:
            try: 
                self.stdout = stdout.read().strip()
            except socket.timeout:
                self.stdout = ""
            try: 
                self.stderr = stderr.read().strip()
            except socket.timeout:
                self.stderr = ""

            if self.stdout:
                self.record(self.stdout)
            if self.stderr:
                self.record(self.stderr)
        finally:
            return self.stdout, self.stderr

    def put(self, local_file, remote_file, msg='file', basename=False):
        self.log("debug", "open sftp and put %s --> %s" % (local_file, remote_file))
        show_name = os.path.basename(remote_file) if basename else remote_file
        self.log('notice', "put %s to %s" % (msg, show_name))
        sftp = self.ssh.open_sftp()
        sftp.put(local_file, remote_file)
        del sftp

    def get(self, remote_file, local_file, msg='file', basename=False):
        self.log('debug', "open sftp and get %s --> %s" % (remote_file, local_file))
        show_name = os.path.basename(local_file) if basename else local_file
        self.log('notice', "get %s to %s" % (msg, show_name))
        sftp = self.ssh.open_sftp()
        sftp.get(remote_file, local_file)
        del sftp

    def kill_task(self, task, signal=0, sleep=0.0, ssh=None):
        prog = re.compile(r"^\S+\s+(\d+) ")
        kill = "kill" if signal == 0 else "kill -%d" % signal
        if ssh:
            ssh_cmd = "ssh %s " % ssh
        else:
            ssh_cmd = ""

        self.send('%sps -ef | grep "%s"' % (ssh_cmd, task))
        findtask = False
        for line in self.stdout.splitlines():
            if "grep" in line:
                continue
            pids = prog.findall(line)
            if pids:
                cmd = "{0}{1} {2}; wait {2} 2>/dev/null".format(ssh_cmd, kill, pids[0])
                self.log("debug", "kill task: %s" % cmd)
                self.send(cmd)
                findtask = True
        if findtask and sleep > 0.0:
            time.sleep(sleep)
        return findtask

